#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
YOLO格式转VOC格式转换脚本
支持将YOLO格式的标注文件转换为VOC格式的XML文件
"""

import os
import xml.etree.ElementTree as ET
from xml.dom import minidom
import cv2
import argparse
from pathlib import Path


class YOLOToVOCConverter:
    def __init__(self, yolo_dir, output_dir, class_names_file=None):
        """
        初始化转换器

        Args:
            yolo_dir: YOLO格式数据集目录
            output_dir: VOC格式输出目录
            class_names_file: 类别名称文件路径
        """
        self.yolo_dir = Path(yolo_dir)
        self.output_dir = Path(output_dir)
        self.class_names = []

        # 创建输出目录
        self.output_dir.mkdir(parents=True, exist_ok=True)
        (self.output_dir / "Annotations").mkdir(exist_ok=True)
        (self.output_dir / "JPEGImages").mkdir(exist_ok=True)
        (self.output_dir / "ImageSets" / "Main").mkdir(parents=True, exist_ok=True)

        # 读取类别名称
        if class_names_file:
            self.load_class_names(class_names_file)

    def load_class_names(self, class_names_file):
        """加载类别名称"""
        try:
            with open(class_names_file, "r", encoding="utf-8") as f:
                self.class_names = [line.strip() for line in f.readlines()]
            print(f"加载了 {len(self.class_names)} 个类别: {self.class_names}")
        except Exception as e:
            print(f"加载类别名称失败: {e}")
            self.class_names = []

    def yolo_to_voc_bbox(self, yolo_bbox, img_width, img_height):
        """
        将YOLO格式边界框转换为VOC格式

        Args:
            yolo_bbox: [center_x, center_y, width, height] (归一化坐标)
            img_width: 图片宽度
            img_height: 图片高度

        Returns:
            (xmin, ymin, xmax, ymax) 像素坐标
        """
        center_x, center_y, width, height = yolo_bbox

        # 转换为像素坐标
        center_x *= img_width
        center_y *= img_height
        width *= img_width
        height *= img_height

        # 计算边界框坐标
        xmin = int(center_x - width / 2)
        ymin = int(center_y - height / 2)
        xmax = int(center_x + width / 2)
        ymax = int(center_y + height / 2)

        # 确保坐标在图片范围内
        xmin = max(0, xmin)
        ymin = max(0, ymin)
        xmax = min(img_width, xmax)
        ymax = min(img_height, ymax)

        return xmin, ymin, xmax, ymax

    def create_voc_xml(self, img_path, img_width, img_height, objects):
        """
        创建VOC格式的XML文件

        Args:
            img_path: 图片路径
            img_width: 图片宽度
            img_height: 图片高度
            objects: 对象列表，每个对象包含 (class_id, bbox)
        """
        # 创建根元素
        annotation = ET.Element("annotation")

        # 添加文件夹信息
        folder = ET.SubElement(annotation, "folder")
        folder.text = "JPEGImages"

        # 添加文件名
        filename = ET.SubElement(annotation, "filename")
        filename.text = Path(img_path).name

        # 添加路径
        path = ET.SubElement(annotation, "path")
        path.text = str(img_path)

        # 添加数据源
        source = ET.SubElement(annotation, "source")
        database = ET.SubElement(source, "database")
        database.text = "Unknown"

        # 添加图片尺寸
        size = ET.SubElement(annotation, "size")
        width_elem = ET.SubElement(size, "width")
        width_elem.text = str(img_width)
        height_elem = ET.SubElement(size, "height")
        height_elem.text = str(img_height)
        depth_elem = ET.SubElement(size, "depth")
        depth_elem.text = "3"

        # 添加分割信息
        segmented = ET.SubElement(annotation, "segmented")
        segmented.text = "0"

        # 添加对象信息
        for class_id, bbox in objects:
            obj = ET.SubElement(annotation, "object")

            # 类别名称
            name = ET.SubElement(obj, "name")
            if class_id < len(self.class_names):
                name.text = self.class_names[class_id]
            else:
                name.text = f"class_{class_id}"

            # 姿态
            pose = ET.SubElement(obj, "pose")
            pose.text = "Unspecified"

            # 截断
            truncated = ET.SubElement(obj, "truncated")
            truncated.text = "0"

            # 困难
            difficult = ET.SubElement(obj, "difficult")
            difficult.text = "0"

            # 边界框
            bndbox = ET.SubElement(obj, "bndbox")
            xmin, ymin, xmax, ymax = bbox

            xmin_elem = ET.SubElement(bndbox, "xmin")
            xmin_elem.text = str(xmin)
            ymin_elem = ET.SubElement(bndbox, "ymin")
            ymin_elem.text = str(ymin)
            xmax_elem = ET.SubElement(bndbox, "xmax")
            xmax_elem.text = str(xmax)
            ymax_elem = ET.SubElement(bndbox, "ymax")
            ymax_elem.text = str(ymax)

        return annotation

    def convert_single_image(self, img_path, txt_path):
        """
        转换单张图片的标注

        Args:
            img_path: 图片路径
            txt_path: 标注文件路径
        """
        # 读取图片获取尺寸
        img = cv2.imread(str(img_path))
        if img is None:
            print(f"无法读取图片: {img_path}")
            return False

        img_height, img_width = img.shape[:2]

        # 读取YOLO标注
        objects = []
        if txt_path.exists():
            with open(txt_path, "r") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        parts = line.split()
                        if len(parts) >= 5:
                            class_id = int(parts[0])
                            bbox = [float(x) for x in parts[1:5]]

                            # 转换为VOC格式
                            voc_bbox = self.yolo_to_voc_bbox(
                                bbox, img_width, img_height
                            )
                            objects.append((class_id, voc_bbox))

        # 创建XML文件
        xml_content = self.create_voc_xml(img_path, img_width, img_height, objects)

        # 保存XML文件
        xml_filename = Path(img_path).stem + ".xml"
        xml_path = self.output_dir / "Annotations" / xml_filename

        # 格式化XML
        rough_string = ET.tostring(xml_content, "unicode")
        reparsed = minidom.parseString(rough_string)
        pretty_xml = reparsed.toprettyxml(indent="  ")

        # 移除空行
        pretty_xml = "\n".join(
            [line for line in pretty_xml.split("\n") if line.strip()]
        )

        with open(xml_path, "w", encoding="utf-8") as f:
            f.write(pretty_xml)

        return True

    def convert_dataset(self, images_dir="obj_train_data", labels_dir="obj_train_data"):
        """
        转换整个数据集

        Args:
            images_dir: 图片目录名
            labels_dir: 标注目录名
        """
        images_path = self.yolo_dir / images_dir
        labels_path = self.yolo_dir / labels_dir

        if not images_path.exists():
            print(f"图片目录不存在: {images_path}")
            return

        # 获取所有图片文件
        image_extensions = {".jpg", ".jpeg", ".png", ".bmp"}
        image_files = []
        for ext in image_extensions:
            image_files.extend(images_path.glob(f"*{ext}"))
            image_files.extend(images_path.glob(f"*{ext.upper()}"))

        print(f"找到 {len(image_files)} 张图片")

        converted_count = 0
        train_list = []

        for img_path in image_files:
            # 对应的标注文件
            txt_path = labels_path / (img_path.stem + ".txt")

            # 转换图片
            if self.convert_single_image(img_path, txt_path):
                # 复制图片到VOC格式目录
                dest_img_path = self.output_dir / "JPEGImages" / img_path.name
                import shutil

                shutil.copy2(img_path, dest_img_path)

                # 添加到训练列表
                train_list.append(img_path.stem)
                converted_count += 1

                if converted_count % 10 == 0:
                    print(f"已转换 {converted_count} 张图片...")

        # 保存训练列表
        train_list_path = self.output_dir / "ImageSets" / "Main" / "train.txt"
        with open(train_list_path, "w") as f:
            for img_name in train_list:
                f.write(f"{img_name}\n")

        print(f"转换完成！共转换 {converted_count} 张图片")
        print(f"输出目录: {self.output_dir}")
        print(f"训练列表: {train_list_path}")


def main():
    parser = argparse.ArgumentParser(description="YOLO格式转VOC格式")
    parser.add_argument("--yolo_dir", required=True, help="YOLO格式数据集目录")
    parser.add_argument("--output_dir", required=True, help="VOC格式输出目录")
    parser.add_argument("--class_names", help="类别名称文件路径")
    parser.add_argument("--images_dir", default="obj_train_data", help="图片目录名")
    parser.add_argument("--labels_dir", default="obj_train_data", help="标注目录名")

    args = parser.parse_args()

    # 创建转换器
    converter = YOLOToVOCConverter(
        yolo_dir=args.yolo_dir,
        output_dir=args.output_dir,
        class_names_file=args.class_names,
    )

    # 执行转换
    converter.convert_dataset(images_dir=args.images_dir, labels_dir=args.labels_dir)


if __name__ == "__main__":
    # 示例用法
    if len(os.sys.argv) == 1:
        # 如果没有命令行参数，使用默认参数
        yolo_dir = r"F:\GitProject\yolov7-pytorch\VOCdevkit\251003\uv75"
        output_dir = r"F:\GitProject\yolov7-pytorch\VOCdevkit\251003\uv75_voc"
        class_names_file = (
            r"F:\GitProject\yolov7-pytorch\VOCdevkit\251003\uv75\obj.names"
        )

        print("使用默认参数进行转换...")
        print(f"YOLO目录: {yolo_dir}")
        print(f"输出目录: {output_dir}")
        print(f"类别文件: {class_names_file}")

        converter = YOLOToVOCConverter(
            yolo_dir=yolo_dir, output_dir=output_dir, class_names_file=class_names_file
        )

        converter.convert_dataset()
    else:
        main()

